"""Utility functions and classes."""
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Callable, List, TypeVar, NamedTuple
from queue import Queue, Empty
import threading
from tqdm import tqdm

T = TypeVar('T')

class OutlineData(NamedTuple):
    """Data needed for outline generation."""
    section_type: str
    paper_text: str
    topic: str
    section_num: int

def parallel_process_with_limit(
    items: List[Any],
    process_func: Callable[[Any], T],
    max_concurrent: int = 20,
    desc: str = None
) -> List[T]:
    """Process items in parallel with limited concurrency.
    
    Args:
        items: List of items to process
        process_func: Function to process each item
        max_concurrent: Maximum number of concurrent threads
        desc: Description for progress tracking
        
    Returns:
        List of processed results in the same order as input items
    """
    results = [None] * len(items)
    task_queue = Queue()
    for i, item in enumerate(items):
        task_queue.put((i, item))
    
    # Create progress bar
    pbar = tqdm(
        total=len(items),
        desc=desc or "Processing",
        unit="item"
    )
    
    # Lock for thread-safe operations
    lock = threading.Lock()
    completed_count = 0
    
    def process_queue():
        nonlocal completed_count
        while True:
            try:
                # Get next task from queue
                index, item = task_queue.get_nowait()
            except Empty:
                # No more tasks in queue
                break
                
            try:
                # Process the item
                result = process_func(item)
                results[index] = result
            except Exception as e:
                print(f"Error processing item {index}: {str(e)}")
            finally:
                # Update progress
                with lock:
                    completed_count += 1
                    pbar.update(1)
                task_queue.task_done()
    
    # Create and start worker threads
    threads = []
    for _ in range(max_concurrent):
        thread = threading.Thread(target=process_queue)
        thread.start()
        threads.append(thread)
    
    # Wait for all tasks to complete
    task_queue.join()
    
    # Wait for all threads to finish
    for thread in threads:
        thread.join()
    
    pbar.close()
    return results


def parallel_outline_generate(
    outline_data_list: List[OutlineData],
    generate_func: Callable[[OutlineData], str],
    desc: str = "Generating outlines"
) -> List[str]:
    """Generate outlines in parallel with limited concurrency."""
    return parallel_process_with_limit(
        items=outline_data_list,
        process_func=generate_func,
        max_concurrent=20,
        desc=desc
    )

def parallel_process(
    items: List[Any],
    process_func: Callable[[Any], T],
    max_workers: int = None,
    desc: str = None
) -> List[T]:
    """Process items in parallel with limited concurrency."""
    return parallel_process_with_limit(
        items=items,
        process_func=process_func,
        max_concurrent=20,
        desc=desc
    ) 